package com.subl.mybatis.aop;

import com.subl.mybatis.util.IdUtils;
import com.subl.mybatis.xa.IssXADataSource;
import com.subl.mybatis.xa.IssXid;
import org.aspectj.lang.annotation.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Lazy;
import org.springframework.jdbc.datasource.ConnectionHolder;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.stereotype.Component;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.DefaultTransactionDefinition;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import javax.sql.DataSource;
import javax.sql.XAConnection;
import javax.transaction.xa.XAException;
import javax.transaction.xa.XAResource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

@Aspect
@Lazy(value = false)
public class GlobalTransactionAspect implements ApplicationContextAware {

    private static final Logger logger = LoggerFactory.getLogger(GlobalTransactionAspect.class);
    private ThreadLocal<TransactionStatus> transactionStatusThreadLocal = new ThreadLocal<>();
    private ThreadLocal<Map<IssXid, XAConnection>> issXidThreadLocal = new ThreadLocal<>();
    private ThreadLocal<Boolean> hasException = new ThreadLocal<>();
    private ApplicationContext applicationContext;

    @Pointcut("@annotation(com.subl.mybatis.annotation.GlobalTransaction)")
    public void cutMethod() {

    }

    @Before("cutMethod()")
    public void doBefore() {
        Map<String, IssXADataSource> beans = applicationContext.getBeansOfType(IssXADataSource.class);
        if (beans.size() == 0) {
            logger.error("未发现XA数据源，非分布式事务场景下不能使用@GlobalTransaction注解！");
            throw new RuntimeException("未发现XA数据源，非分布式事务场景下不能使用@GlobalTransaction注解！");
        }
        //初始化线程异常状态
        hasException.set(false);
        //只连接一个数据库使用本地事务
        if (beans.size() == 1) {
            logger.warn("只有一个XA数据源，不会产生分布式事务，将使用本地事务处理");
            DataSourceTransactionManager transactionManager = applicationContext.getBean(DataSourceTransactionManager.class);
            DefaultTransactionDefinition def = new DefaultTransactionDefinition(TransactionDefinition.PROPAGATION_REQUIRES_NEW);
            TransactionStatus status = transactionManager.getTransaction(def);
            transactionStatusThreadLocal.set(status);
            return;
        }
        int i = 0;
        Map<IssXid, XAConnection> issXidMap = new HashMap<>();
        Iterator<Map.Entry<String, IssXADataSource>> iterator = beans.entrySet().iterator();
        while (iterator.hasNext()) {
            try {
                IssXADataSource xaDataSource = iterator.next().getValue();
                XAConnection xaConnection = xaDataSource.getXAConnection();
                XAResource xaResource = xaConnection.getXAResource();
                IssXid xid = new IssXid(IdUtils.genGtrid().getBytes(), ("bqual" + i).getBytes(), 1);
                xaResource.start(xid, XAResource.TMNOFLAGS);
                issXidMap.put(xid, xaConnection);
                //在ThreadLocal中绑定connection
                bindResource(xaConnection.getConnection(), xaDataSource.getDataSource());
            } catch (SQLException e) {
                throw new RuntimeException(e);
            } catch (XAException e) {
                throw new RuntimeException(e);
            }
            i++;
        }
        issXidThreadLocal.set(issXidMap);

    }

    @After("cutMethod()")
    public void doAfter() {
        Map<String, IssXADataSource> beans = applicationContext.getBeansOfType(IssXADataSource.class);
        //只连接一个数据库使用本地事务，在此进行事务提交
        if (beans.size() == 1) {

            DataSourceTransactionManager transactionManager = applicationContext.getBean(DataSourceTransactionManager.class);
            TransactionStatus status = transactionStatusThreadLocal.get();
            //判断是否发生异常
            Boolean flag = hasException.get();
            if (flag) {
                //发生异常时回滚事务
                logger.info("发生异常，使用本地事务回滚处理");
                transactionManager.rollback(status);
            } else {
                logger.info("使用本地事务提交");
                transactionManager.commit(status);
            }
            transactionStatusThreadLocal.remove();
            return;
        }
        Map<IssXid, XAConnection> issXidMap = issXidThreadLocal.get();
        Iterator<Map.Entry<IssXid, XAConnection>> iterator = issXidMap.entrySet().iterator();
        // 各分支事务执行完毕
        while (iterator.hasNext()) {
            Map.Entry<IssXid, XAConnection> e = iterator.next();
            IssXid xid = e.getKey();
            try {
                XAResource xaResource = e.getValue().getXAResource();
                xaResource.end(xid, XAResource.TMSUCCESS);
            } catch (Exception ex) {
                logger.error("XA事务end阶段发生异常", ex);
                throw new RuntimeException(ex);
            }
        }
        // 判断业务逻辑执行过程中是否发生异常
        Boolean flag = hasException.get();
        if (flag) {
            // 回滚XA事务
            xaRollback();
            closeXaConnection();
            //清除ThreadLocal中的状态信息
            issXidThreadLocal.remove();
            hasException.remove();
            return;
        }
        try {
            // 两阶段提交协议第一阶段
            iterator = issXidMap.entrySet().iterator();
            while (iterator.hasNext()) {

                try {
                    Map.Entry<IssXid, XAConnection> entry = iterator.next();
                    XAResource xaResource = entry.getValue().getXAResource();
                    IssXid xid = entry.getKey();
                    xaResource.prepare(xid);
                } catch (Exception e) {
                    logger.error("XA事务prepare阶段发生异常", e);
                    throw e;
                }
            }
            // 两阶段提交协议第二阶段
            iterator = issXidMap.entrySet().iterator();
            while (iterator.hasNext()) {
                try {
                    Map.Entry<IssXid, XAConnection> entry = iterator.next();
                    XAResource xaResource = entry.getValue().getXAResource();
                    IssXid xid = entry.getKey();
                    xaResource.commit(xid, false);
                } catch (Exception e) {
                    logger.error("XA事务commit阶段发生异常", e);
                    throw e;
                }
            }

        } catch (Exception e) {
            xaRollback();
            throw new RuntimeException(e);
        } finally {
            closeXaConnection();
            //清除ThreadLocal中的状态信息
            issXidThreadLocal.remove();
            hasException.remove();


        }
    }

    @AfterThrowing("cutMethod()")
    public void afterThrowing() {
        // 执行业务逻辑过程中发生异常，需要回滚事务
        hasException.set(true);
    }

    private void closeXaConnection() {
        Map<IssXid, XAConnection> issXidMap = issXidThreadLocal.get();
        // 关闭连接
        Iterator<Map.Entry<IssXid, XAConnection>> iterator = issXidMap.entrySet().iterator();
        while (iterator.hasNext()) {
            try {
                Map.Entry<IssXid, XAConnection> entry = iterator.next();
                XAConnection xaConnection = entry.getValue();
                xaConnection.close();
            } catch (Exception e) {
                logger.error("XA事务关闭连接发生异常", e);
            }
        }
    }

    private void xaRollback() {
        // 执行业务逻辑过程中发生异常，需要回滚事务
        logger.info("执行业务逻辑过程中发生异常，需要回滚事务");
        Map<IssXid, XAConnection> issXidMap = issXidThreadLocal.get();
        Iterator<Map.Entry<IssXid, XAConnection>> iterator = issXidMap.entrySet().iterator();
        // 各分支事务执行完毕
        while (iterator.hasNext()) {
            Map.Entry<IssXid, XAConnection> e = iterator.next();
            IssXid xid = e.getKey();
            try {
                XAResource xaResource = e.getValue().getXAResource();
                xaResource.rollback(xid);
            } catch (Exception ex) {
                logger.error("XA事务rollback阶段发生异常", ex);
                throw new RuntimeException(ex);
            }
        }
    }


    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    private void bindResource(Connection con, DataSource dataSource) {
        ConnectionHolder holderToUse = new ConnectionHolder(con);
        holderToUse.requested();
        holderToUse.setSynchronizedWithTransaction(true);
        TransactionSynchronizationManager.bindResource(dataSource, holderToUse);
    }

}
